import threading
import queue
import time
from typing import Callable, Dict, List, Literal, Optional, Union
import logging
import uuid

# from abc import ABC, abstractmethod
logger = logging.getLogger(__name__)


class Task:
    """
    表示一个可执行的任务对象。
    """

    def __init__(self, task_id: str = "", *args, **kwargs):
        self.task_id = task_id or str(uuid.uuid4())
        self.created_at_ms: int = int(time.time() * 1000)
        self.finished_at_ms: int = -1
        self.log_buffer: List[str] = []
        self.status: Literal["waiting", "running", "finished"] = "waiting"
        self.result: Optional[Union[Dict, List]] = None
        self.user_id: int = (
            -1
        )  # -1 indicates a task that is not associated with any user.

    def get_display_name(self) -> str:
        """
        Return the display name of the task
        """
        return f"{self.__class__.__name__}"

    def get_status_html(self) -> List[str]:
        """
        Return the status information in HTML format
        """
        return ["<p>Normal</p>"]

    def run(self):
        """
        执行任务，并捕获输出到日志缓冲区。
        如果用户未覆盖此方法，则抛出 NotImplementedError。
        """
        raise NotImplementedError("必须覆盖 run 方法以定义任务逻辑")

    def write_log(self, message: str):
        """
        写入日志到当前任务的日志缓冲区。
        """
        self.log_buffer.append(message)

    def finish(self):
        """
        完成任务，并更新任务的状态。
        """
        self.status = "finished"
        self.finished_at_ms = int(time.time() * 1000)


class TaskManager:
    def __init__(self, pool_size: int):
        if pool_size <= 0:
            raise ValueError("pool_size 必须大于 0")

        self.task_queue = queue.Queue()
        self.current_tasks: dict[str, Task] = {}  # task_id --> task
        self.threads = []
        self.is_running = True
        self.lock = threading.Lock()

        for _ in range(pool_size):
            thread = threading.Thread(target=self._worker)
            thread.daemon = True
            self.threads.append(thread)
            thread.start()

    def _worker(self):
        while self.is_running:
            task: Task = self.task_queue.get()

            if task is None:
                continue
            self.current_tasks[task.task_id] = task
            try:
                task.status = "running"
                task.run()
                task.finish()
            except Exception as e:
                import traceback

                traceback.print_exc()
                logger.error(f"Error occurred when executing ${task}: {e}")
            finally:
                self.task_queue.task_done()

    def submit(self, task: Task):
        """
        提交一个任务到线程池中。
        """
        self.task_queue.put(task)

    def shutdown(self):
        """
        关闭线程池，等待所有任务完成。
        """
        self.is_running = False
        for _ in self.threads:
            self.task_queue.put(None)  # 唤醒线程退出循环
        for thread in self.threads:
            thread.join()

    def get_create_container_task_status(self, task_id: str) -> dict:
        """
        获取指定 task_id 的 CreateContainerTask 状态和 HTML 日志。

        :param task_id: 任务的唯一标识符。
        :return: 包含任务状态和 HTML 日志的字典。
        """
        with self.lock:
            if task_id in self.current_tasks:
                task = self.current_tasks[task_id]
                # if isinstance(task, CreateContainerTask):
                return {
                    "status": task.status,
                    "status_html": task.get_status_html(),
                    "result": task.result,
                }
                # else:
                #     raise ValueError("指定的任务不是 CreateContainerTask 类型")
            else:
                raise KeyError(f"任务 {task_id} 不存在")

    def get_task_logs(self, task_id: str) -> List[str]:
        """
        获取指定任务的日志，并清空已获取的日志。

        :param task_id: 任务的唯一标识符。
        :return: 日志列表。
        """
        with self.lock:
            if task_id in self.current_tasks:
                task = self.current_tasks[task_id]
                items = task.log_buffer[:]
                task.log_buffer = []
                return items
            else:
                return []
